import os
import time
import torch
import torchvision
from torchvision import datasets, models, transforms
from torch.utils.data import DataLoader
from pycocotools.coco import COCO
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
from torch import nn, optim

# -------------------------------
# 辅助工具：日志记录器
# -------------------------------
class MetricLogger:
    def __init__(self, delimiter="\t"):
        self.delimiter = delimiter

    def update(self, **kwargs):
        print(self.delimiter.join(f"{k}: {v:.4f}" for k, v in kwargs.items()))

    def log_every(self, iterable, print_freq, header=None):
        i = 0
        if header is not None:
            print(header)
        for obj in iterable:
            yield obj
            i += 1
            if i % print_freq == 0:
                print(f"Iteration {i}")

# -------------------------------
# 辅助函数：标注格式转换
# 将COCO的 [x, y, width, height] 转换为 [x_min, y_min, x_max, y_max]
# -------------------------------
def convert_target(target):
    boxes = []
    labels = []
    for obj in target:
        x, y, w, h = obj['bbox']
        boxes.append([x, y, x + w, y + h])
        labels.append(obj['category_id'])
    boxes = torch.tensor(boxes, dtype=torch.float32)
    labels = torch.tensor(labels, dtype=torch.int64)
    return {"boxes": boxes, "labels": labels}

# -------------------------------
# 辅助函数：计算IoU（交并比）
# -------------------------------
def compute_iou(boxA, boxB):
    xA = max(boxA[0], boxB[0])
    yA = max(boxA[1], boxB[1])
    xB = min(boxA[2], boxB[2])
    yB = min(boxA[3], boxB[3])
    interArea = max(0, xB - xA) * max(0, yB - yA)
    boxAArea = (boxA[2] - boxA[0]) * (boxA[3] - boxA[1])
    boxBArea = (boxB[2] - boxB[0]) * (boxB[3] - boxB[1])
    iou = interArea / float(boxAArea + boxBArea - interArea + 1e-6)
    return iou

# -------------------------------
# 辅助函数：计算准确率（简化版本）
# 对每个真实目标，若存在预测框（分数>=0.5）与其类别一致且IoU>=0.5，则认为检测正确
# -------------------------------
def evaluate_accuracy(model, data_loader, device, max_batches=None):
    model.eval()
    correct = 0
    total = 0
    batch_count = 0
    for images, targets in data_loader:
        if max_batches is not None and batch_count >= max_batches:
            break
        images = [img.to(device) for img in images]
        # 转换标注格式
        targets = [convert_target(t) for t in targets]
        targets = [{k: v.to(device) for k, v in t.items()} for t in targets]

        with torch.no_grad():
            outputs = model(images)
        for output, target in zip(outputs, targets):
            gt_boxes = target["boxes"]
            gt_labels = target["labels"]
            total += len(gt_boxes)
            pred_boxes = output["boxes"]
            pred_labels = output["labels"]
            pred_scores = output["scores"]
            for gt_box, gt_label in zip(gt_boxes, gt_labels):
                found = False
                for pb, pl, ps in zip(pred_boxes, pred_labels, pred_scores):
                    if ps < 0.5:
                        continue
                    if pl == gt_label:
                        iou_val = compute_iou(gt_box.cpu().numpy(), pb.cpu().numpy())
                        if iou_val >= 0.5:
                            found = True
                            break
                if found:
                    correct += 1
        batch_count += 1
    return correct / total if total > 0 else 0.0

# -------------------------------
# 数据预处理与数据加载
# -------------------------------
from torch.utils.data import DataLoader
import torchvision.transforms as T

def collate_fn(batch):
    images, targets = zip(*batch)
    # 直接堆叠已经是 tensor 的图像
    images = torch.stack(images, dim=0)
    # 对目标标注转换为 tensor，如果需要的话
    targets = [{k: torch.tensor(v) for k, v in t.items()} for t in targets]
    return images, targets


# 使用自定义 collate_fn
def get_data_loaders(data_dir, batch_size=4, num_workers=4, image_size=(800, 800)):
    transform = T.Compose([
        T.Resize(image_size),
        T.ToTensor(),
        T.Normalize(mean=[0.485, 0.456, 0.406],
                    std=[0.229, 0.224, 0.225]),
    ])

    train_data = datasets.CocoDetection(root=os.path.join(data_dir, 'train2014'),
                                          annFile=os.path.join(data_dir, 'instances_train2014.json'),
                                          transform=transform)

    val_data = datasets.CocoDetection(root=os.path.join(data_dir, 'val2014'),
                                        annFile=os.path.join(data_dir, 'instances_val2014.json'),
                                        transform=transform)

    train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True, num_workers=num_workers, collate_fn=collate_fn)
    val_loader = DataLoader(val_data, batch_size=batch_size, shuffle=False, num_workers=num_workers, collate_fn=collate_fn)
    return train_loader, val_loader, transform



# -------------------------------
# 加载预训练模型并进行微调
# -------------------------------
def get_model(num_classes=91):
    model = models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    model.roi_heads.box_predictor = models.detection.faster_rcnn.FastRCNNPredictor(in_features, num_classes)
    return model

# -------------------------------
# 修改训练函数：返回平均损失值
# -------------------------------
def train_one_epoch(model, optimizer, data_loader, device, epoch):
    model.train()
    metric_logger = MetricLogger(delimiter="  ")
    header = f"Epoch #{epoch}"
    epoch_loss = 0.0
    num_batches = 0
    for images, targets in metric_logger.log_every(data_loader, 100, header):
        images = [image.to(device) for image in images]
        targets = [convert_target(t) for t in targets]
        targets = [{k: v.to(device) for k, v in t.items()} for t in targets]

        loss_dict = model(images, targets)
        losses = sum(loss for loss in loss_dict.values())

        optimizer.zero_grad()
        losses.backward()
        optimizer.step()

        epoch_loss += losses.item()
        num_batches += 1
        metric_logger.update(loss=losses.item())
    return epoch_loss / num_batches if num_batches > 0 else 0.0

# -------------------------------
# 推理与结果可视化函数
# -------------------------------
def plot_result(image, outputs):
    plt.imshow(image)
    boxes = outputs['boxes'].cpu().numpy()
    labels = outputs['labels'].cpu().numpy()
    scores = outputs['scores'].cpu().numpy()

    for box, label, score in zip(boxes, labels, scores):
        if score > 0.5:  # 阈值
            plt.gca().add_patch(plt.Rectangle((box[0], box[1]),
                                              box[2] - box[0],
                                              box[3] - box[1],
                                              linewidth=2, edgecolor='r', facecolor='none'))
            plt.text(box[0], box[1], f"{label}: {score:.2f}",
                     color='red', fontsize=12)
    plt.show()

# -------------------------------
# 主函数
# -------------------------------
if __name__ == '__main__':
    # 设置数据集路径
    data_dir = "data/coco2014/"  # 请替换为本地COCO数据集路径

    # 获取数据加载器与预处理函数
    train_loader, val_loader, transform = get_data_loaders(data_dir, batch_size=4, num_workers=4)

    # 初始化模型并设置设备
    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
    model = get_model(num_classes=91)
    model.to(device)

    # 设置优化器
    params = [p for p in model.parameters() if p.requires_grad]
    optimizer = optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005)

    # 训练参数及记录列表
    num_epochs = 10
    train_loss_list = []
    train_acc_list = []
    val_acc_list = []

    for epoch in range(num_epochs):
        print(f"----- Epoch {epoch} -----")
        train_loss = train_one_epoch(model, optimizer, train_loader, device, epoch)
        # 为加快评估速度，这里仅在部分batch上计算准确率
        train_acc = evaluate_accuracy(model, train_loader, device, max_batches=50)
        val_acc = evaluate_accuracy(model, val_loader, device, max_batches=20)
        train_loss_list.append(train_loss)
        train_acc_list.append(train_acc)
        val_acc_list.append(val_acc)
        print(f"Epoch {epoch} -- Loss: {train_loss:.4f}, Training Accuracy: {train_acc:.4f}, Validation Accuracy: {val_acc:.4f}")

    # 保存模型
    torch.save(model.state_dict(), 'fasterrcnn_coco.pth')

    # 绘制训练过程图像：损失及准确率变化
    epochs = range(num_epochs)
    plt.figure(figsize=(15, 5))

    plt.subplot(1, 3, 1)
    plt.plot(epochs, train_loss_list, marker='o')
    plt.title("训练损失变化")
    plt.xlabel("Epoch")
    plt.ylabel("Loss")

    plt.subplot(1, 3, 2)
    plt.plot(epochs, train_acc_list, marker='o')
    plt.title("训练准确率变化")
    plt.xlabel("Epoch")
    plt.ylabel("Accuracy")

    plt.subplot(1, 3, 3)
    plt.plot(epochs, val_acc_list, marker='o')
    plt.title("验证准确率变化")
    plt.xlabel("Epoch")
    plt.ylabel("Accuracy")

    plt.tight_layout()
    plt.show()

    # 示例推理：加载一张图像进行检测并展示结果
    image_path = "path/to/some/image.jpg"  # 请替换为实际图像路径
    image = Image.open(image_path)
    image_tensor = transform(image).unsqueeze(0).to(device)

    model.eval()
    with torch.no_grad():
        prediction = model(image_tensor)

    plot_result(np.array(image), prediction[0])
